from abc import ABC, abstractmethod
from typing import List, Dict, Any

class SchedulerBase(ABC):
    """Abstract base class for task schedulers.
    
    This class defines the interface that all scheduler implementations must follow.
    Schedulers are responsible for task distribution and synchronization in the distributed system.
    """
    
    def __init__(self, config: Dict[str, Any]):
        """Initialize the scheduler with configuration.
        
        Args:
            config: Dictionary containing scheduler configuration parameters
        """
        self.config = config
        self.workers = []
        self.tasks = []
        self.sync_points = []
    
    @abstractmethod
    def initialize(self) -> bool:
        """Initialize the scheduler with the current system state.
        
        Returns:
            bool: True if initialization is successful, False otherwise
        """
        pass
    
    @abstractmethod
    def assign_tasks(self) -> List[Dict[str, Any]]:
        """Assign tasks to available workers based on the scheduling strategy.
        
        Returns:
            List[Dict[str, Any]]: List of task assignments, each containing task and worker information
        """
        pass
    
    @abstractmethod
    def update_status(self, worker_id: str, task_id: str, status: str) -> None:
        """Update the status of a task on a specific worker.
        
        Args:
            worker_id: Identifier of the worker
            task_id: Identifier of the task
            status: New status of the task
        """
        pass
    
    @abstractmethod
    def handle_failure(self, worker_id: str, task_id: str) -> None:
        """Handle task failure on a worker.
        
        Args:
            worker_id: Identifier of the failed worker
            task_id: Identifier of the failed task
        """
        pass
    
    @abstractmethod
    def get_sync_points(self) -> List[int]:
        """Get the list of synchronization points in the computation graph.
        
        Returns:
            List[int]: List of layer indices where synchronization is required
        """
        pass
    
    @abstractmethod
    def free_nodes(self, completed_tasks: List[str]) -> List[str]:
        """Identify nodes that can be freed after task completion.
        
        Args:
            completed_tasks: List of completed task identifiers
            
        Returns:
            List[str]: List of node identifiers that can be freed
        """
        pass